# -*- coding: utf-8 -*-
"""
Created on Tue Sep 26 14:58:01 2017

@author: pyh
"""
import os
import shutil
import sys
import numpy as np
sys.path.append('/public/home/users/ruc001/bin/aBest')
import vasp

NP=sys.argv[1]
PBS_NODEFILE=sys.argv[2]
num1=int(sys.argv[3])
nvtstep=1999
cosang=1.0
sinang=0.0
direction=-1
displace=5
dpv=10
mode='z'


run='/public/software/mpi/openmpi/1.6.5/intel/bin/mpirun -np '+NP+' -machinefile '+PBS_NODEFILE+' --mca btl self,sm,openib --bind-to-core  ~/bin/vasp5c  | tee sout'

try:
    os.mkdir('NVT')
    os.chdir('NVT')
    os.symlink('../INCAR_NVT','./INCAR')
    os.symlink('../KPOINTS','./KPOINTS')
    os.symlink('../POTCAR','./POTCAR')
    os.symlink('../POSCAR','./POSCAR')
    if os.path.exists('vdw_kernel.bindat'):
        os.symlink('../vdw_kernel.bindat','./'+ind[:]+'/vdw_kernel.bindat')
    os.system(run)
    os.chdir('../')
except:
    pass

if not (os.path.exists('NVEmin') or os.path.exists('NVEmax')):    
    os.chdir('NVT')
    potim,nsw,natom,lc,posa=vasp.readvasprun()

    vela=vasp.velcal(lc,posa,potim)
    velas=[[] for i in range(nvtstep,len(vela))]

    for i in range(nvtstep,len(vela)):
        velas[i-nvtstep]=vela[i][num1-1]

    fvmax,fvmin=vasp.findlimit(velas)

    vasp.writepos('POSCAR','pos_vmax',posa[nvtstep+fvmax],vela[nvtstep+fvmax])
    vasp.writepos('POSCAR','pos_vmin',posa[nvtstep+fvmin],vela[nvtstep+fvmin])

    shutil.copy('pos_vmax','../pos_vmax')
    shutil.copy('pos_vmin','../pos_vmin')

    os.chdir('../')


try:
    os.mkdir('NVEmin')
    os.chdir('NVEmin')
    os.symlink('../INCAR_NVE','./INCAR')
    os.symlink('../KPOINTS','./KPOINTS')
    os.symlink('../POTCAR','./POTCAR')
    if os.path.exists('vdw_kernel.bindat'):
        os.symlink('../vdw_kernel.bindat','./'+ind[:]+'/vdw_kernel.bindat')
    shutil.copy('../pos_vmax','./pos_vmax')
    shutil.copy('../record','./record')
    os.chdir('..')
except:
    pass

os.chdir('NVEmin')

fid = open('record','rt')
line=fid.readlines()
fid.close()

try:
    vel_start=float(line[-1].replace('\n',''))
except:
    vel_start=0

ap,avel=vasp.readpos('pos_vmax')
avel[num1-1,2]=avel[num1-1,2]-vel_start*cosang
avel[num1-1,int(abs(direction))-1]=avel[num1-1,int(abs(direction))-1]+direction/abs(direction)*vel_start*sinang
vasp.writepos('pos_vmax','POSCAR',ap,avel)


while 'end\n' not in line:
    vasp.filemv('INCAR')
    ind=line[-1].replace('\n','')
    os.chdir(ind)
    shutil.copy('../record','./record')
    os.system(run)

    potim,nsw,natom,lc,posa=vasp.readvasprun()
    vela1=vasp.velcal(lc,posa,potim,num1)
    check1=vasp.checkaway(lc,posa,vela1,num1,True,displace,dpv,mode)

    ns=vasp.bisec(check1)

    if ns:
        ap,avel=vasp.readpos('POSCAR')
        avel[num1-1,2]=avel[num1-1,2]-ns*cosang
        avel[num1-1,int(abs(direction))-1]=avel[num1-1,int(abs(direction))-1]+direction/abs(direction)*ns*sinang
        vasp.writepos('POSCAR','pos_next',ap,avel)
    try:
        shutil.copy('./pos_next','../POSCAR')
    except:
        pass
    shutil.copy('./record','../record')

    fid = open('record','rt')
    line=fid.readlines()
    fid.close()

    os.chdir('../')

os.chdir('../')

try:
    os.mkdir('NVEmax')
    os.chdir('NVEmax')
    os.symlink('../INCAR_NVE','./INCAR')
    os.symlink('../KPOINTS','./KPOINTS')
    os.symlink('../POTCAR','./POTCAR')
    if os.path.exists('vdw_kernel.bindat'):
        os.symlink('../vdw_kernel.bindat','./'+ind[:]+'/vdw_kernel.bindat')
    shutil.copy('../pos_vmin','./pos_vmin')
    shutil.copy('../record','./record')
    os.chdir('..')
except:
    pass

os.chdir('NVEmax')

fid = open('record','rt')
line=fid.readlines()
fid.close()

try:
    vel_start=float(line[-1].replace('\n',''))
except:
    vel_start=0

ap,avel=vasp.readpos('pos_vmin')
avel[num1-1,2]=avel[num1-1,2]-vel_start*cosang
avel[num1-1,int(abs(direction))-1]=avel[num1-1,int(abs(direction))-1]+direction/abs(direction)*vel_start*sinang
vasp.writepos('pos_vmin','POSCAR',ap,avel)


while 'end\n' not in line:
    vasp.filemv('INCAR')
    ind=line[-1].replace('\n','')
    os.chdir(ind)
    shutil.copy('../record','./record')
    os.system(run)

    potim,nsw,natom,lc,posa=vasp.readvasprun()
    vela1=vasp.velcal(lc,posa,potim,num1)
    check1=vasp.checkaway(lc,posa,vela1,num1,True,displace,dpv,mode)

    ns=vasp.bisec(check1)

    if ns:
        ap,avel=vasp.readpos('POSCAR')
        avel[num1-1,2]=avel[num1-1,2]-ns*cosang
        avel[num1-1,int(abs(direction))-1]=avel[num1-1,int(abs(direction))-1]+direction/abs(direction)*ns*sinang
        vasp.writepos('POSCAR','pos_next',ap,avel)

    try:
        shutil.copy('./pos_next','../POSCAR')
    except:
        pass

    shutil.copy('./record','../record')

    fid = open('record','rt')
    line=fid.readlines()
    fid.close()

    os.chdir('../')